#!/user/bin/env python3
# -*- coding: utf-8 -*-
import cv2
import numpy as np
import os


def calculate_brightness(img):
    """计算图片亮度状态"""
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    height, width = gray_img.shape
    size = gray_img.size

    # 计算直方图
    hist = cv2.calcHist([gray_img], [0], None, [256], [0, 256])

    # 计算偏离均值(128)的程度
    reduce_matrix = np.full((height, width), 128)
    shift_value = gray_img - reduce_matrix
    shift_sum = np.sum(shift_value)
    da = shift_sum / size

    # 计算平均偏差
    ma = np.sum(np.abs(np.arange(256) - 128 - da) * hist.flatten())
    m = np.abs(ma / size)

    # 亮度系数
    k = np.abs(da) / m
    
    return k, da


def process_image(image_path, output_dir):
    """处理单张图片"""
    img = cv2.imread(image_path)
    if img is None:
        print(f"无法读取图片: {image_path}")
        return

    k, da = calculate_brightness(img)

    # 判断亮度状态
    if k > 1:
        if da > 0:
            status = "Overexposed"
            color = (0, 0, 255)  # 红色表示过亮
        else:
            if da >= -80:
                status = "Slightly Dark"
                color = (255, 165, 0)  # 橙色
            elif -110 <= da < -80:
                status = "Moderately Dark"
                color = (0, 0, 255)  # 红色
            else:
                status = "Very Dark"
                color = (255, 0, 0)  # 蓝色
    else:
        status = "Normal"
        color = (0, 255, 0)  # 绿色表示正常

    # 在图片上渲染状态文本和参数
    text1 = f"Brightness: {status}"
    text2 = f"k = {k:.4f}, da = {da:.2f}"

    cv2.putText(img, text1, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
    cv2.putText(img, text2, (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

    # 在图片上渲染状态文本
    text = f"Brightness: {status}"
    cv2.putText(img, text, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)

    # 保存处理后的图片
    filename = os.path.basename(image_path)
    output_path = os.path.join(output_dir, filename)
    cv2.imwrite(output_path, img)
    print(f"已处理: {filename} -> {status}")


def process_folder(input_dir, output_dir):
    """处理文件夹中的所有图片"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
            image_path = os.path.join(input_dir, filename)
            process_image(image_path, output_dir)


if __name__ == "__main__":
    input_folder = "/ai_data1/wireope_seg/testData/DA5804532"  # 输入文件夹路径
    output_folder = "output_images"  # 输出文件夹路径



    process_folder(input_folder, output_folder)
    print("所有图片处理完成！")